import torch.utils.data as Data
import torch
import torchvision
from joblib import dump, load
import torch.nn as nn
import os
import PIL
import matplotlib.pyplot as plt
from sklearn.metrics import confusion_matrix
from sklearn.metrics import classification_report
from sklearn.manifold import TSNE
import seaborn as sns
import numpy as np
import warnings
import csv
from matplotlib import MatplotlibDeprecationWarning
from sklearn.decomposition import PCA
warnings.filterwarnings('ignore', category=MatplotlibDeprecationWarning)

torch.manual_seed(100) # 设置随机种子，以使实验结果具有可重复性
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 模型 测试集 测试
def model_test(model, test_loader, batch_size):
    true_labels = [] # 存储类别标签
    predicted_labels = [] # 存储预测标签

    # 提取特征
    test_original_features = [] # 原始测试集 特征
    test_original_labels= [] # 原始测试集 标签
    test_model_features = [] # 测试集经过训练后的模型得到的特征

    model = model.to(device)

    # 定义损失函数和优化函数
    loss_function = nn.CrossEntropyLoss(reduction='sum')
    # 每一个epoch后，在验证集上验证实验结果
    with torch.no_grad():
        loss_test = 0.
        correct_test = 0
        # print(len(test_loader.dataset))
        for j, (data, label) in enumerate(test_loader):
            # print(data.view(batch_size,-1).shape)
            # test_original_features += data.view(batch_size,-1).tolist() # 提取特征
            test_original_features += data.view(data.shape[0],-1).tolist() # 提取特征
            # print(data.view(batch_size, -1).shape)
            test_original_labels += label.tolist() # 提取标签

            # 将模型者设置为评估模型
            model.eval()

            data, label = data.to(device), label.to(device)
            pre = model(data)

            test_model_features += pre.tolist() # 提取特征
            correct_test += torch.sum(pre.argmax(dim=1).view(-1)==label.view(-1)).item()
            loss = loss_function(pre, label)
            loss_test += loss.item()
            # 计算个数
            predicted = torch.argmax(pre, dim=1)
            true_labels.extend(label.tolist())
            predicted_labels.extend(predicted.tolist())
        print(f'test_loss:{loss_test/len(test_loader)}, test_Acc:{correct_test/len(test_loader.dataset)}%')
        return true_labels, predicted_labels, test_original_features, test_original_labels, test_model_features

# 原始 测试集  进行t-SNE降维
def plot_original_TSNE(test_original_features, test_original_labels, num_classes, class_labels, savepath):
    features = np.array(test_original_features)
    labels = np.array(test_original_labels)
    # 初始化t-SNE模型
    tsne = TSNE(n_components=2, random_state=100)
    # print(type(features))
    # 对原始数据进行t-SNE降维
    features_tsne = tsne.fit_transform(features)
    colors = ['#4C9AC9', '#E2745E', '#91AD5A', '#F9E9A4']

    # 可视化结果
    plt.figure(figsize=(10,8), dpi=100)
    # 绘制每个类别的散点图，并指定label
    for i in range(num_classes):
        plt.scatter(features_tsne[labels == i,0],features_tsne[labels==i,1], c=colors[i],label=class_labels[i], alpha=0.8)
    plt.title('Original-Test t-SNE Visualization')
    # plt.xlabel('tSNE Dimension 1')
    # plt.ylabel('tSNE Dimension 2')
    # 显示legend
    plt.legend()
    # plt.show()
    plt.savefig(os.path.join(savepath,'rawdata_TSNE.png'), dpi=100)
    plt.close()
def plot_original_3D(test_original_features, test_original_labels, num_classes, class_labels, savepath):
    features = np.array(test_original_features)
    labels = np.array(test_original_labels)
    # 初始化t-SNE模型
    # tsne = TSNE(n_components=2, random_state=42)
    method = TSNE(n_components=3, random_state=100)
    # method = PCA(n_components=3)
    features_tsne = method.fit_transform(features)
    # print(type(features))
    # 对原始数据进行t-SNE降维
    # features_tsne = tsne.fit_transform(features)
    colors = ['#4C9AC9', '#E2745E', '#91AD5A', '#F9E9A4']
    # 可视化结果
    plt.figure(figsize=(10,8), dpi=100)
    ax = plt.axes(projection='3d')
    # 绘制每个类别的散点图，并指定label
    for i in range(num_classes):
        ax.scatter3D(features_tsne[labels == i,0],features_tsne[labels==i,1], c=colors[i],label=class_labels[i], alpha=0.6)
    plt.title('Original-Test t-SNE Visualization')
    # plt.xlabel('tSNE Dimension 1')
    # plt.ylabel('tSNE Dimension 2')
    # 显示legend
    plt.legend()
    # plt.show()
    plt.savefig(os.path.join(savepath,'rawdata_3D.png'), dpi=100)
    plt.close()
# 训练好的模型 测试集进行t-SNE降维
def plot_model_TSNE(test_model_features, test_original_labels, num_classes, class_labels, savepath):
    features = np.array(test_model_features)
    labels = np.array(test_original_labels)

    # 初始化t-SNE降维
    tsne = TSNE(n_components=2, random_state=100)
    TSNE()
    # 可视化结果
    features_tsne = tsne.fit_transform(features)
    colors = ['#4C9AC9', '#E2745E', '#91AD5A', '#F9E9A4']
    # 绘制每个类别的散点图，并致电label
    for i in range(num_classes):
        plt.scatter(features_tsne[labels==i, 0], features_tsne[labels==i,1], c=colors[i],label=class_labels[i], alpha=0.8)
    plt.title('Model-Test t-SNE Visualization')
    # plt.xlabel('tSNE Dimension 1')
    # plt.ylabel('tSNE Dimension 2')
    # 显示legend
    plt.legend()
    # plt.show()
    plt.savefig(os.path.join(savepath, 'modeldata_TSNE.png'), dpi=300)
    plt.close()
def plot_model_3D(test_model_features, test_original_labels, num_classes, class_labels, savepath):
    features = np.array(test_model_features)
    labels = np.array(test_original_labels)

    # 初始化t-SNE降维
    # tsne = TSNE(n_components=2, random_state=42)
    method = TSNE(n_components=3, random_state=100)
    # method = PCA(n_components=3)
    features_tsne = method.fit_transform(features)
    plt.figure(figsize=(10,6), dpi=100)
    ax = plt.axes(projection='3d')
    # # 可视化结果
    # features_tsne = tsne.fit_transform(features)
    colors = ['#4C9AC9','#E2745E','#91AD5A','#F9E9A4']
    # 绘制每个类别的散点图，并致电label
    for i in range(num_classes):
        ax.scatter3D(features_tsne[labels==i, 0], features_tsne[labels==i,1],c=colors[i], label=class_labels[i], alpha=0.6)
    plt.title('Model-Test t-SNE Visualization')
    # plt.xlabel('tSNE Dimension 1')
    # plt.ylabel('tSNE Dimension 2')
    # 显示legend
    plt.legend()
    # plt.show()
    plt.savefig(os.path.join(savepath, 'modeldata_3D.png'), dpi=300)
    plt.close()
# 标签真实值和预测值对比
def plot_true_pre_compare(class_labels, predicted_labels,savepath):
    # 类别标签
    x = [i for i in range(len(class_labels))]
    # 可视化结果
    plt.figure(figsize=(10,6), dpi=100)
    # 绘制真实标签的折线图
    plt.scatter(x, predicted_labels, color='green', label='Predicted Labels')

    plt.title('Compare labels real and predict values')
    plt.xlabel('samples')
    plt.ylabel('labels values')
    plt.xticks(class_labels)
    plt.legend()
    plt.grid(True)
    # plt.show()
    plt.savefig(os.path.join(savepath, 'compare_samples.png'), dpi=300)

if __name__=='__main__':
    # 加载模型
    param_toolboxPath = r'D:\Project_mb\modelParameters\BCI42a\StarNet'
    result_toolboxPath = r'D:\Project_mb\CWTResult\BCI42a\StarNet'
    subjects = [f'sub0{i + 1}' for i in range(9)]
    for i in range(len(subjects)):
        data_dir = os.path.join(result_toolboxPath, subjects[i])
        model_dir = os.path.join(param_toolboxPath, subjects[i])
        model = torch.load(os.path.join(model_dir, 'CWT_best_model.pt'), map_location=torch.device('cpu'))

        # 加载测试集
        test_loader = load(os.path.join(data_dir,'val_loader.pkl'))
        # print(len(test_loader.dataset))
        batch_size = 32
        # 模型测试
        true_labels, predicted_labels, test_original_features, tes_original_labels, test_model_features = model_test(model, test_loader, batch_size)
        # print(np.array(test_original_features).shape)
        # print(np.array(test_original_features).shape)
        # print(np.array(tes_original_labels).shape)
        # 计算每一类的分类准确率
        report = classification_report(true_labels, predicted_labels, digits=4)
        print(report)
        with open(os.path.join(result_toolboxPath,'output.csv'),mode='a',newline='') as file:
            csv_writer = csv.writer(file)
            csv_writer.writerow(report)
        file.close()
        # 混淆矩阵
        # 原始标签和自定义标签的映射
        label_maping = {0: 'left hand', 1:' right hand', 2:'tongue', 3:'foot'}
        confusion_mat = confusion_matrix(true_labels, predicted_labels)

        # 绘制混淆矩阵
        plt.figure(figsize=(10, 8), dpi=300)
        sns.heatmap(confusion_mat, xticklabels=label_maping.values(), yticklabels=label_maping.values(), annot=True,
                    fmt='d', cmap='summer')
        plt.xlabel('Predicted Labels')
        plt.ylabel('True Labels')
        plt.title('Confusion Matrix')
        # plt.show()
        plt.savefig(os.path.join(data_dir, 'Confusion_Matrix.png'), dpi=300)
        plt.close()

        num_classes = 4
        # 类别标签
        class_labels = ['left hand', 'right hand', 'foot', 'tongue']
        # 原始测试集进行t-SNE降维
        # print(np.array(len(test_original_features)))
        # print(np.array(test_original_features).shape)
        plot_original_TSNE(test_original_features, tes_original_labels, num_classes, class_labels,data_dir)
        plot_original_3D(test_original_features, tes_original_labels, num_classes, class_labels,data_dir)
        # # 训练好的模型测试集进行t-SNE降维
        plot_model_TSNE(test_model_features, tes_original_labels, num_classes, class_labels, data_dir)
        plot_model_3D(test_model_features, tes_original_labels, num_classes, class_labels, data_dir)
        # 绘制标签真实值和预测值对比
        # plot_true_pre_compare(true_labels, predicted_labels, data_dir)
        print(f'{subjects[i]} is done!')
